from __future__ import division
from functools import partial
from progress import ProgressDisplay
from osm_utils import *
from problem_generator import *
from map_problem import *
from search_runners import *
from plot_map import *
from cool import *

def Examine1(map, results):
    ''' first graph - comparing run times '''
    from matplotlib import pyplot as plt

    calc_gamma_f(map)
    x = []
    y_c_runtime = []
    y_f_runtime = []
    y_s_runtime = []
    gamma_f = calc_gamma_f(map)
    for result in results:
        initial, goal, c_nodes, c_runtime, c_cost, f_nodes, f_runtime, f_cost, s_nodes, s_runtime, s_cost = result
        C_f_star = f_cost # beacuse the huristic func H_f() is admissible
        x.append(gamma_f * C_f_star)
        y_c_runtime.append(c_runtime)
        y_f_runtime.append(f_runtime)
        y_s_runtime.append(s_runtime)

    values = zip(x, y_c_runtime, y_f_runtime, y_s_runtime)
    values.sort(key = lambda v : v[0])
    x = [v[0] for v in values]
    y_c_runtime = [v[1] for v in values]
    y_f_runtime = [v[2] for v in values]
    y_s_runtime = [v[3] for v in values]

    plt.plot(x, y_c_runtime, 'r', label = '$C_c$')
    plt.plot(x, y_f_runtime, 'g', label = '$C_f$')
    plt.plot(x, y_s_runtime, 'b', label = '$C_s$')
    
    plt.legend(loc = 0)
    plt.title('Experiment 1', fontsize=20)
    plt.ylabel('run time [s]',fontsize=16)
    plt.xlabel(r'optimal route in time $\gamma_f\cdot C_f$',fontsize=16)
    plt.show()


def Examine2(map, results):
    ''' 2nd graph - optimal route distance'''
    from matplotlib import pyplot as plt

    calc_gamma_f(map)
    x = []
    y_c_cost = []
    gamma_f = calc_gamma_f(map)
    for result in results:
        initial, goal, c_nodes, c_runtime, c_cost, f_nodes, f_runtime, f_cost, s_nodes, s_runtime, s_cost = result
        C_f_star = f_cost # beacuse the huristic func H_f() is admissible
        x.append(gamma_f * C_f_star)
        y_c_cost.append(c_cost)

    values = zip(x, y_c_cost)
    values.sort(key = lambda v : v[0])
    x = [v[0] for v in values]
    y_c_cost = [v[1] for v in values]

    plt.plot(x, y_c_cost, 'r')
    
    plt.title('Experiment 2', fontsize=20)
    plt.ylabel('optimal route in length',fontsize=16)
    plt.xlabel(r'optimal route in time ($\gamma_f\cdot C_f$)',fontsize=16)
    plt.show()



def Examine3(map, results):
    ''' 3rd graph - Sigma '''
    from matplotlib import pyplot as plt

    calc_gamma_f(map)
    x = []
    y_s_cost = []
    gamma_f = calc_gamma_f(map)
    for result in results:
        initial, goal, c_nodes, c_runtime, c_cost, f_nodes, f_runtime, f_cost, s_nodes, s_runtime, s_cost = result
        C_f_star = f_cost # beacuse the huristic func H_f() is admissible
        x.append(gamma_f * C_f_star)
        y_s_cost.append(s_cost)

    values = zip(x, y_s_cost)
    values.sort(key = lambda v : v[0])
    x = [v[0] for v in values]
    y_s_cost = [v[1] for v in values]

    plt.plot(x, y_s_cost, 'r')
    
    plt.title('Experiment 3', fontsize=20)
    plt.ylabel('$C_\Sigma$',fontsize=16)
    plt.xlabel(r'optimal route in time ($\gamma_f\cdot C_f$)',fontsize=16)
    plt.show()


def ExamineFuncsCompare(map, initial_key, goal_key):
    ''' same route with different cost funcs '''
    plt.title('Experiment funcs route comparison', fontsize=20)
    
    plt.subplot(2, 2, 1)
    node, developed_nodes_count, run_time, path_cost  = RunAstar(map, initial_key, goal_key, C_c, H_c)
    plot_path(map, (n.action for n in node.path() if n.action is not None))
    node, developed_nodes_count, run_time, path_cost  = RunAstar(map, goal_key, initial_key, C_c, H_c)
    plot_path(map, (n.action for n in node.path() if n.action is not None))
    plt.title('$C_s$')    
    
    plt.subplot(2, 2, 2)
    node, developed_nodes_count, run_time, path_cost  = RunAstar(map, initial_key, goal_key, C_f, H_f)
    plot_path(map, (n.action for n in node.path() if n.action is not None))
    node, developed_nodes_count, run_time, path_cost  = RunAstar(map, goal_key, initial_key, C_f, H_f)
    plot_path(map, (n.action for n in node.path() if n.action is not None))
    plt.title('$C_f$')
    
    plt.subplot(2, 2, 3)
    node, developed_nodes_count, run_time, path_cost  = RunSigmaAstar(map, initial_key, goal_key, (0.4, 0.4))
    plot_path(map, (n.action for n in node.path() if n.action is not None))
    node, developed_nodes_count, run_time, path_cost  = RunSigmaAstar(map, goal_key, initial_key, (0.4, 0.4))
    plot_path(map, (n.action for n in node.path() if n.action is not None))
    plt.title('$C_\Sigma$')

    plt.subplot(2, 2, 4)
    #node, developed_nodes_count, run_time, path_cost = RunCoolPathSearch(map, initial_key, goal_key)
    import cPickle as pickle
    try:
        node = pickle.load(open('cool.pkl','rb'))
        plot_path(map, (Link(n.action.source[1], n.action.target[1], 0, 0, 0) for n in node.path() if n.action is not None))
    except IOError:
        print 'ERROR! please run cool.py in PyPy to create the path'
    plt.title('Cool path')

    plt.show()


if __name__ == '__main__':
    map = CountryMap().LoadMap()
    results1 = LoadfFromFile('output1.csv', float)
    Examine1(map, results1)
    Examine2(map, results1)
    Examine3(map, results1)
    ExamineFuncsCompare(map, 108838, 111600)